#include <avl_bst.h>
#include <stdlib.h>

static void fill_new_node(bst_t* bst_p, avl_node_t* avl_node_p,
			  const void* data_p) {
	avl_node_p->height = 0;
	bst_node_t* node_p = &(avl_node_p->node);
	node_p->left_pn = nullptr;
	node_p->right_pn = nullptr;
	bst_cpy_node_data(bst_p, node_p, data_p);
}

static iword_s avl_get_height(const avl_node_t* node_pn) {
	if (node_pn == nullptr) {
		return -1;
	} else {
		return node_pn->height;
	}
}

static iword_s avl_get_bst_height(bst_node_t* node_pn) {
	avl_node_t* avl_node_p = avl_get_node_pn(node_pn);
	return avl_get_height(avl_node_p);
}

static void modify_ll(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	bst_node_t* left_p = parent_p->left_pn;
	parent_p->left_pn = left_p->right_pn;
	left_p->right_pn = parent_p;
	*parent_link_pp = left_p;
	avl_node_t* parent_avl_p = avl_get_node_pn(parent_p);
	parent_avl_p->height -= 1;
}

static void modify_lr(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	bst_node_t* left_p = parent_p->left_pn;
	bst_node_t* left_right_p = left_p->right_pn;
	parent_p->left_pn = left_right_p->right_pn;
	left_p->right_pn = left_right_p->left_pn;
	left_right_p->left_pn = left_p;
	left_right_p->right_pn = parent_p;
	*parent_link_pp = left_right_p;
	avl_node_t* parent_avl_p = avl_get_node_pn(parent_p);
	uword_s height = parent_avl_p->height;
	parent_avl_p->height = height - 1;
	avl_node_t* left_avl_p = avl_get_node_pn(left_p);
	left_avl_p->height = height - 1;
	avl_node_t* left_right_avl_p = avl_get_node_pn(left_right_p);
	left_right_avl_p->height = height;
}

static void modify_l(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	iword_s left_left_height =
		avl_get_bst_height(parent_p->left_pn->left_pn);
	iword_s left_right_height =
		avl_get_bst_height(parent_p->left_pn->right_pn);
	if (left_left_height > left_right_height) {
		modify_ll(parent_link_pp, parent_p);
	} else {
		modify_lr(parent_link_pp, parent_p);
	}
}

static void modify_rl(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	bst_node_t* right_p = parent_p->right_pn;
	bst_node_t* right_left_p = right_p->left_pn;
	parent_p->right_pn = right_left_p->left_pn;
	right_p->left_pn = right_left_p->right_pn;
	right_left_p->right_pn = right_p;
	right_left_p->left_pn = parent_p;
	*parent_link_pp = right_left_p;
	avl_node_t* parent_avl_p = avl_get_node_pn(parent_p);
	uword_s height = parent_avl_p->height;
	parent_avl_p->height = height - 1;
	avl_node_t* right_avl_p = avl_get_node_pn(right_p);
	right_avl_p->height = height - 1;
	avl_node_t* right_left_avl_p = avl_get_node_pn(right_left_p);
	right_left_avl_p->height = height;
}

static void modify_rr(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	bst_node_t* right_p = parent_p->right_pn;
	parent_p->right_pn = right_p->left_pn;
	right_p->left_pn = parent_p;
	*parent_link_pp = right_p;
	avl_node_t* parent_avl_p = avl_get_node_pn(parent_p);
	parent_avl_p->height -= 1;
}

static void modify_r(bst_node_t** parent_link_pp, bst_node_t* parent_p) {
	iword_s right_left_height =
		avl_get_bst_height(parent_p->right_pn->left_pn);
	iword_s right_right_height =
		avl_get_bst_height(parent_p->right_pn->right_pn);
	if (right_left_height > right_right_height) {
		modify_rl(parent_link_pp, parent_p);
	} else {
		modify_rr(parent_link_pp, parent_p);
	}
}

// 进入此函数时，parent->node的高度未更新，其所有子结点的高度已更新。
static bool_word_s balance_one_layer(bst_link_t* parent_link_p) {
	bst_node_t** parent_link_pp = parent_link_p->node_link_pnp;
	bst_node_t* parent_p = parent_link_p->node_pn;

	avl_node_t* parent_avl_p = avl_get_node_pn(parent_p);
	iword_s left_height = avl_get_bst_height(parent_p->left_pn);
	iword_s right_height = avl_get_bst_height(parent_p->right_pn);
	iword_s old_height = parent_avl_p->height;

	// 如果parent的高度未变（只要大于两子树高度必然未变），直接退出balance。
	iword_s max_child_height;
	if (left_height > right_height) {
		max_child_height = left_height;
	} else {
		max_child_height = right_height;
	}
	if (old_height > max_child_height) {
		return false;
	}
	// 注意：如果parent的高度变了，不应立即更新高度，而是先balance。
	if (unlikely(left_height > right_height + 1)) {
		modify_l(parent_link_pp, parent_p);
	} else if (unlikely(right_height > left_height + 1)) {
		modify_r(parent_link_pp, parent_p);
	}
	// 如果不需要balance，需要更新parent的高度，并继续向上检查。
	else {
		parent_avl_p->height = max_child_height + 1;
		return true;
	}

	// 如果经过了balance，则后续不再需要balance。
	return false;
}

// 进入此函数时，parent->node的高度未更新，其所有子结点的高度已更新。
static void modify_balance(bst_link_t link_a[],
			   bst_link_t* init_parent_link_p) {
	bst_link_t* parent_link_p = init_parent_link_p;
	while (1) {
		bool_word_s need_continue = balance_one_layer(parent_link_p);
		if (!need_continue) {
			break;
		} else if (parent_link_p > link_a) {
			parent_link_p--;
		} else {
			break;
		}
	}
}

bbst_errno_s avl_add(bst_t* bst_p, const void* data_p) {
	bst_link_t space_a[MAX_HEIGHT];
	bst_link_t* space_p = space_a;
	bst_link_t result =
		bst_search_link(bst_p, data_p, &space_p, space_a + MAX_HEIGHT);
	avl_node_t* node_pn = avl_get_node_pn(result.node_pn);

	// 如果data已存在，直接返回。
	if (node_pn != nullptr) {
		return BBST_DATA_EXIST;
	}

	bst_node_t** link_pnp = result.node_link_pnp;
	avl_node_t* new_node_pn =
		malloc(sizeof(avl_node_t) + bst_p->data_size);
	if (unlikely(new_node_pn == nullptr)) {
		return BBST_MEM_ALLOC_FAIL;
	}
	avl_node_t* new_avl_node_p = new_node_pn;
	fill_new_node(bst_p, new_avl_node_p, data_p);
	bst_node_t* new_bst_node_p = &(new_avl_node_p->node);
	*link_pnp = new_bst_node_p;

	if (likely(space_p >= space_a + 2)) {
		// space_a数组的有效部分的node_pn不可能为空，可以看作node_p。
		avl_node_t* parent_p = avl_get_node_pn(space_p[-2].node_pn);

		// 新插入结点的父结点高度必然为1。
		parent_p->height = 1;
	}

	// 太矮的树不可能失衡，必须有二级父结点。
	if (likely(space_p >= space_a + 3)) {
		space_p[-1].node_pn = new_bst_node_p;

		// 至此，space_pna数组直到space_pnp（不含）区间都非空，
		// 可视作space_pa。
		modify_balance(space_a, space_p - 3);
	}
	return BBST_SUCCESS;
}
